package net.tp.parallel;

import java.util.*;
import java.util.function.*;
import java.util.concurrent.*;
import static net.tp.util.Utils.*;

public class ParallelRadixSort {

    public static void sort(int[] a) {
        sort(a, 4, Runtime.getRuntime().availableProcessors(), true);
    }

    public static void sort(int[] a, int d, int K, boolean asc) {
        int N = a.length;
        int M = 1 + (N - 1)/K; // subarray's size for each thread
        int[] b = new int[a.length];
        int[] src = a; int[] dest = b;

        ExecutorService service = Executors.newFixedThreadPool(K);

        for (int i = 0; i < 32; i += d) {
            radixPass(src, dest, i, d, K, asc, service);
            int[] tmp = src; src = dest; dest = tmp;
        }
        radixPass(src, dest, 31, 1, K, !asc, service); // fix the sign
        if (dest != a) parallelCopy(dest, 0, a, 0, N);

        service.shutdownNow();
    }

    static void radixPass(int[] src, int[] dest, int pos, int d, int K, boolean asc, ExecutorService service) {
        final int N = src.length;
        final int M = 1 + (N - 1) / K; // subarray's size for each thread
        final int R = 1 << d;
        final int mask = R-1;
        final int[][] cnt = new int[K+1][R];

        parallelRanges(service, 0, N, K, (k, from, to) -> {
            for (int i = from; i < to; i++)
                cnt[k][ (src[i]>>>pos)&mask ]++;
        });

        for (int r = 0; r < R; r++)
            for (int sum = 0, k = 0; k < K+1; k++) {
                int tmp = cnt[k][r];
                cnt[k][r] = sum;
                sum += tmp;
            }

        if (asc) {
            for (int r = 1; r < R; r++)
                cnt[K][r] += cnt[K][r-1];
            for (int r = 1; r < R; r++)
                for (int k = 0; k < K; k++)
                    cnt[k][r] += cnt[K][r-1];
        }
        else {
            for (int r = R-2; r >= 0; r--)
                cnt[K][r] += cnt[K][r+1];
            for (int r = R-2; r >= 0; r--)
                for (int k = 0; k < K; k++)
                    cnt[k][r] += cnt[K][r+1];
        }

        parallelRanges(service, 0, N, K, (k, from, to) -> {
            for (int i = from; i < to; i++)
                dest[cnt[k][ (src[i]>>>pos)&mask ]++] = src[i];
        });
    }

}
